package com.sunday.common.trace.zipkin.study.spring.thread;

import brave.Span;
import brave.Tracer;
import brave.Tracing;
import brave.propagation.CurrentTraceContext;
import brave.propagation.TraceContext;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.Optional;
import java.util.concurrent.FutureTask;
import java.util.concurrent.ThreadPoolExecutor;

import static brave.Span.Kind.CONSUMER;

@Slf4j
@Configuration
public class ThreadPoolConfig {

    @Bean
    public ThreadPoolTaskExecutor threadPool() {

        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        executor.setThreadNamePrefix("ThreadPool-");
        //核心线程数
        executor.setCorePoolSize(5);
        //最大线程数
        executor.setMaxPoolSize(10);
        //队列最大长度 >=mainExecutor.maxSize
        executor.setQueueCapacity(1);
        //线程池维护线程所允许的空闲时间
        executor.setKeepAliveSeconds(60);
        //线程池对拒绝任务(无线程可用)的处理策略
        executor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());
        executor.initialize();
        return executor;
    }


    @Bean
    public ThreadPoolTaskExecutor threadPool2() {

        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        executor.setThreadNamePrefix("ThreadPool2-");
        //核心线程数
        executor.setCorePoolSize(5);
        //最大线程数
        executor.setMaxPoolSize(10);
        //队列最大长度 >=mainExecutor.maxSize
        executor.setQueueCapacity(1);
        //线程池维护线程所允许的空闲时间
        executor.setKeepAliveSeconds(60);
        //线程池对拒绝任务(无线程可用)的处理策略
        executor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());
        executor.initialize();

        Tracing tracing = Tracing.current();
        Tracer tracer = tracing.tracer();
        CurrentTraceContext currentTraceContext = tracing.currentTraceContext();

        /**
         *  ScopedSpan startScopedSpan
         *  如果有{@link #currentSpan()}，返回一个新的子span;如果没有，返回一个新的trace。的
         *  结果是“当前跨度”，直到{@link ScopedSpan#finish()}被调用。
         */
//        executor.setTaskDecorator(task -> () -> {
//            ScopedSpan scopedSpan = tracer.startScopedSpan(Optional.ofNullable(task.getClass())
//                    .filter(clazz -> clazz.isAssignableFrom(FutureTask.class))
//                    .map(clazz -> executor.getThreadNamePrefix() + "FutureTask")
//                    .orElseGet(() -> executor.getThreadNamePrefix() + "Runnable"));
//            try {
//                task.run();
//            } finally {
//                scopedSpan.finish();
//            }
//        });

        executor.setTaskDecorator(task -> {
            TraceContext context = currentTraceContext.get();
            Span span = tracer.newChild(context);
            return () -> {
                try (Tracer.SpanInScope scope = tracer.withSpanInScope(span)) {
                    span.kind(CONSUMER);
                    span.name(Optional.ofNullable(task.getClass())
                            .filter(clazz -> clazz.isAssignableFrom(FutureTask.class))
                            .map(f -> executor.getThreadNamePrefix() + "submit")
                            .orElseGet(() -> executor.getThreadNamePrefix() + "execute"));
                    span.start();
                    task.run();
                } finally {
                    span.finish();
                }
            };
        });

//        executor.setTaskDecorator(task -> {
//            TraceContext context = currentTraceContext.get();
//            return () -> {
//                try (CurrentTraceContext.Scope scope = currentTraceContext.maybeScope(context)) {
//                    task.run();
//                }
//            };
//        });

        return executor;
    }

}
